from stock import Neeq, TencentStock, Snowball, My_tushare
from pymongo import MongoClient
import pymongo
import arrow
from tqdm import tqdm
import time
from multiprocessing.dummy import Process

Mongo = MongoClient('192.168.1.220',29001)
NEEQDB = Mongo['stock_neeq_v1']
TencentDB = Mongo['stock_tencent_v1']
SnowBallDB = Mongo['stock_snowball_v1']
TushareDB = Mongo['tushare_v1']
datestr = arrow.now().format('YYYYMMDD')

def neeq(version):
    # one day one version
    neeq = Neeq()

    if NEEQDB['index'].find_one({'version':version}):
        symbols = [symbol.get('symbol') for symbol in NEEQDB['index'].find({'version':version},{'symbol':1})]
    else:
        coms = neeq.get_neeq_coms()

        for com in coms:
            his_id = com['_id']
            com['_id'] = com['symbol']
            com['version'] = version

            NEEQDB['index'].delete_one({'_id':com['_id']})
            NEEQDB['index'].insert_one(com)
            
            com['_id'] = his_id
            NEEQDB['index_history'].insert_one(com)

        symbols = [symbol.get('symbol') for symbol in coms]

    for symbol in tqdm(symbols, desc='get_neeq_detail'):    
        if not NEEQDB['detail'].find_one({'symbol':symbol,'version':version}):
            com = neeq.get_neeq_detail(symbol)

            his_id = com['_id']
            com['_id'] = com['symbol']
            com['version'] = version

            NEEQDB['detail'].delete_one({'_id':com['_id']})
            NEEQDB['detail'].insert_one(com)
            
            com['_id'] = his_id
            NEEQDB['detail_history'].insert_one(com)


def Tencent(version):


    symbols = [symbol.get('symbol') for symbol in NEEQDB['index'].find({'version':version},{'symbol':1})]
    finish = [symbol.get('symbol') for symbol in TencentDB['detail'].find({'version':version},{'symbol':1})]
    symbols = set(symbols) - set(finish)
    ts = TencentStock()
    for symbol in tqdm(symbols,desc='get_basic_info'):
        com = ts.get_basic_info(symbol, 'nq')

        his_id = com['_id']
        com['_id'] = com['symbol']
        com['version'] = version

        TencentDB['detail'].delete_one({'_id':com['_id']})
        TencentDB['detail'].insert_one(com)
        
        com['_id'] = his_id
        TencentDB['detail_history'].insert_one(com)


def snowball(version):


    symbols = [symbol.get('symbol') for symbol in NEEQDB['index'].find({'version':version},{'symbol':1})]
    finish = [symbol.get('symbol') for symbol in SnowBallDB['detail'].find({'version':version},{'symbol':1})]
    detail_symbols = set(symbols) - set(finish)

    sb = Snowball()
    for symbol in tqdm(detail_symbols, desc='get_basic_info'):
        com = sb.get_basic_info(symbol, 'oc')
        if com:
            his_id = com['_id']
            com['_id'] = com['symbol']
            com['version'] = version

            SnowBallDB['detail'].delete_one({'_id':com['_id']})
            SnowBallDB['detail'].insert_one(com)
            
            com['_id'] = his_id
            SnowBallDB['detail_history'].insert_one(com)

        time.sleep(0.5)

    finish = [symbol.get('symbol') for symbol in SnowBallDB['capital_structure'].find({'version':version},{'symbol':1})]
    cs_symbols = set(symbols) - set(finish)
    for symbol in tqdm(cs_symbols,desc='get_capital_structure'):
        data = sb.get_capital_structure(symbol, 'oc')
        try:
            if data:
                SnowBallDB['capital_structure'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass
        time.sleep(0.5)

    finish = [symbol.get('symbol') for symbol in SnowBallDB['financial_indicators'].find({'version':version},{'symbol':1})]
    fi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(fi_symbols,desc='get_financial_indicators'):
        data = sb.get_financial_indicators(symbol, 'oc')
        try:
            if data:
                SnowBallDB['financial_indicators'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass
        time.sleep(0.5)

    finish = [symbol.get('symbol') for symbol in SnowBallDB['history_stock'].find({'version':version},{'symbol':1})]
    hi_symbols = set(symbols) - set(finish)
    for symbol in tqdm(hi_symbols,desc='get_h_stock'):
        data = sb.get_h_stock(symbol, 'oc')
        try:
            if data:
                SnowBallDB['history_stock'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass
        time.sleep(0.5)

def tushare(version):


    mt = My_tushare()
    sb = Snowball()
    data_area = mt.get_my_area_classified()
    # try:
    #     TushareDB['area_classified'].insert_many(data_area,ordered=False)

    # except pymongo.errors.BulkWriteError:
    #     # repeat insert
    #     pass

    for com in data_area:
        try:
            his_id = com['_id']
            com['_id'] = com['symbol']
            com['version'] = version

            TushareDB['area_classified'].delete_one({'_id':com['_id']})
            TushareDB['area_classified'].insert_one(com)
            
            com['_id'] = his_id
            TushareDB['area_classified_history'].insert_one(com)
        except pymongo.errors.DuplicateKeyError:
            pass



    for each in tqdm(data_area,desc='stock like OC'):
        symbol = each.get('symbol')
        typ = 'SH' if symbol.startswith('6') else 'SZ'

        finish = set(sbl.get('symbol') for sbl in SnowBallDB['detail'].find({'version':version},{'symbol':1}))
        if symbol not in finish:

            com = sb.get_basic_info(symbol, typ)
            try:
                if com:
                    his_id = com['_id']
                    com['_id'] = com['symbol']
                    com['version'] = version

                    SnowBallDB['detail'].delete_one({'_id':com['_id']})
                    SnowBallDB['detail'].insert_one(com)
                    
                    com['_id'] = his_id
                    SnowBallDB['detail_history'].insert_one(com)
            except pymongo.errors.DuplicateKeyError:
                pass

        time.sleep(0.2)

        data = sb.get_capital_structure(symbol, typ)
        try:
            if data:
                SnowBallDB['capital_structure'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            pass
        time.sleep(0.2)

        data = sb.get_financial_indicators(symbol, typ)
        try:
            if data:
                SnowBallDB['financial_indicators'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            pass
        time.sleep(0.2)

        data = sb.get_h_stock(symbol, typ)
        try:
            if data:
                SnowBallDB['history_stock'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            pass
        time.sleep(0.2)

    for each in tqdm(data_area,desc='get_my_h_data'):
        symbol = each.get('symbol')
        data = mt.get_my_h_data(symbol)
        try:
            if data:
                TushareDB['history_data'].insert_many(data,ordered=False)
        except pymongo.errors.BulkWriteError:
            # repeat insert
            pass


def job_1():
    neeq()
    ps = [Process(target=fun) for fun in [Tencent,snowball]]
    list(map(lambda x:x.start(),ps))
    list(map(lambda x:x.join(),ps))

def job_2():
    tushare()

# def job_1():
#   print(1)
#   time.sleep(2)

# def job_2():
#   print(2)
#   time.sleep(4)

if __name__=='__main__':
    # ps = [Process(target=fun) for fun in [job_1,job_2]]
    # list(map(lambda x:x.start(),ps))
    # list(map(lambda x:x.join(),ps))
    # # ps[0].start()
    # # ps[0].join()
    # print(3)
    # # test()
    versions = NEEQDB['index'].distinct('version')
    if not versions:
        versions = [0]
    max_version = max(versions)
    if NEEQDB['index'].find_one({'update':{'$gte':arrow.now().floor('day').datetime},'version':max_version}):
        version = max_version
    else:
        version = max_version + 1

    print(version)
    version = 3
    neeq(version)
    tushare(version)